# -*- coding: utf-8 -*-
from sklearn.model_selection import train_test_split
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import jieba

model_file = 'my_FC.pth'

def vectorize_sentence(sentence, vector_table):
    words = sentence.split()  # 假设以空格分隔
    sentence_vectors = []

    for word in words:
        if word in vector_table:
            sentence_vectors.append(vector_table[word])
        else:
            sentence_vectors.append(np.zeros(300))  # 如果词不在表中，使用零向量

    # 合并句子向量（例如取平均）
    return np.mean(sentence_vectors, axis=0)

class SimpleNN(nn.Module):
    def __init__(self, input_dim):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_dim, 128)  # 第一层，输入300维，输出128维
        self.fc2 = nn.Linear(128, 64)  # 第二层，输出64维
        self.fc3 = nn.Linear(64, 2)  # 输出层，2个类别（0或1）

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


model = SimpleNN(input_dim=300)
criterion = nn.CrossEntropyLoss()  # 使用交叉熵损失
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
if os.path.exists(model_file):
    model.load_state_dict(torch.load(model_file))
    print("模型权重已加载。")
else:
    print("模型文件不存在，无法加载。")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)  # 将模型移动到GPU

vector_table = {}
with open('../data/sgns.zhihu.bigram', 'r', encoding='utf-8') as f:
    for line in f:
        parts = line.strip().split()
        word = parts[0]
        vector = np.array(parts[1:], dtype=np.float32)  # 转换为numpy数组
        vector_table[word] = vector
print("词典已加载。")


while True:
    sentence = input("请输入句子\n")
    if sentence:
        sentence = " ".join(jieba.cut(sentence, cut_all=False, HMM=True))
        print("sentence:", sentence)
        sentence_vector = vectorize_sentence(sentence, vector_table)
        print(sentence_vector[:10])
        # Convert the sentence vector into a PyTorch tensor
        sentence_tensor = torch.tensor(sentence_vector, dtype=torch.float32).to(device)
        result = model(sentence_tensor)
        predicted_class = torch.argmax(result).item()
        if predicted_class == 1:
            print("结果: 正性")
        else:
            print("结果: 负性")
